//
// Created by limo on 2025/9/26.
//

#pragma once
#include <iostream>
#include <boost/asio.hpp>
#include "tcp_connection.h"

// 导入 boost::asio
using boost::asio::ip::tcp;

namespace net {

    class TcpServer {
    public:
        using ConnectionPtr = Connection::Ptr;
        // Server 接收到新连接的回调
        using NewConnectionHandler = std::function<void(ConnectionPtr)>;

        // 构造函数：接受端口号和线程数量
        TcpServer(unsigned short port, size_t numThreads = std::thread::hardware_concurrency())
                : m_nextIoContext(0),
                  m_acceptorContext(),
                  m_acceptor(m_acceptorContext.get_executor(), tcp::endpoint(tcp::v4(), port)),
                  m_numThreads(numThreads) {

            // 初始化 io_context 列表
            if (numThreads == 0) numThreads = 1; // 至少一个线程
            for (size_t i = 0; i < numThreads; ++i) {
                m_ioContexts.emplace_back(std::make_shared<boost::asio::io_context>());
                // 使用 work_guard 确保 io_context 在没有工作时也不会退出，直到显式调用 stop()
                m_workGuards.emplace_back(boost::asio::make_work_guard(m_ioContexts.back()->get_executor()));
            }
        }

        // 析构函数：确保清理工作
        ~TcpServer() {
            stop();
        }

        // 设置新连接创建时的回调
        void setNewConnectionHandler(NewConnectionHandler handler) {
            m_newConnectionHandler = handler;
        }

        void start() {
            doAccept();

            // 启动工作线程池
            for (size_t i = 0; i < m_numThreads; ++i) {
                // 确保每个线程都运行自己的 io_context
                m_ioThreads.emplace_back([this, i] {
                    std::cout << "Thread " << i << " started running io_context." << std::endl;
                    m_ioContexts[i]->run();
                    std::cout << "Thread " << i << " finished." << std::endl;
                });
            }

            std::cout << "Acceptor Thread started running context." << std::endl;
            boost::asio::executor_work_guard<boost::asio::io_context::executor_type> acceptor_work(m_acceptorContext.get_executor());
            m_acceptorContext.run();

            // 当 run() 返回后，等待 Worker 线程全部退出
            for (auto& t : m_ioThreads) {
                if (t.joinable()) {
                    t.join();
                }
            }
        }

        void stop() {
            // 1. 停止 acceptor context
            m_acceptorContext.stop();

            // 2. 停止所有 worker contexts
            for (auto& ctx : m_ioContexts) {
                ctx->stop();
            }
        }

    private:
        // **独立的 Acceptor Context**
        boost::asio::io_context m_acceptorContext;
        tcp::acceptor m_acceptor;

        // **用于 I/O 任务的 Context Pool**
        std::vector<std::shared_ptr<boost::asio::io_context>> m_ioContexts;
        std::vector<boost::asio::executor_work_guard<boost::asio::io_context::executor_type>> m_workGuards;
        std::vector<std::thread> m_ioThreads;

        // 轮询（Round Robin）机制分配 io_context
        size_t m_nextIoContext;
        size_t m_numThreads;
        NewConnectionHandler m_newConnectionHandler;

        // 获取下一个用于新连接的 io_context
        boost::asio::io_context& getNextIoContext() {
            boost::asio::io_context& ctx = *m_ioContexts[m_nextIoContext];
            m_nextIoContext = (m_nextIoContext + 1) % m_numThreads;
            return ctx;
        }

        void doAccept() {
            // 创建一个新的 TcpConnection 对象，并将其绑定到下一个可用的 io_context
            auto& nextCtx = getNextIoContext();
            // 创建一个新的 TcpConnection 对象来承载新连接
            auto newConnection = std::make_shared<TcpConnection>(nextCtx);
            m_acceptor.async_accept(
                    newConnection->socket(),
                    [this, newConnection](boost::system::error_code ec) {
                        if (!ec) {
                            // 成功接受连接
                            if (m_newConnectionHandler) {
                                // 触发新连接回调，让用户设置该连接上的各种回调
                                m_newConnectionHandler(newConnection);
                            }
                            // 启动新连接的消息循环
                            newConnection->start();
                        }
                        // 无论成功与否，继续等待下一个连接
                        doAccept();
                    });
        }
    };

} // net
